import asynchat, asyncore, socket, struct, heapq, time, core

nodes = core.nodes
hooks = {}

class Channel(asynchat.async_chat):
    def __init__(self, sock = None):
        asynchat.async_chat.__init__(self, sock)

    def init_accept(self):
        self.set_terminator(2)
        self.stage = 0
        self.peer = None
        self.hook = None
        self.buf = ''

    def init_connect(self, remote, listen):
        self.create_socket(socket.AF_INET, socket.SOCK_STREAM)
        self.connect(remote)
        self.push(struct.pack('H', listen))
        self.set_terminator(2)
        self.stage = 1
        self.peer = core.addr_to_ident(remote)
        self.hook = hooks.get(self.peer, None)
        nodes[self.peer] = self
        self.buf = ''

    def collect_incoming_data(self, data):
        self.buf += data

    def found_terminator(self):
        if self.stage == 2:
            self.stage = 1
            self.set_terminator(2)
            if self.hook:
                self.hook(self.buf, self.peer)
#           else:
#               try:
#                   hooks[None](self.buf, self.peer)
#               except KeyError:
#                   pass
        elif self.stage == 1:
            self.stage = 2
            self.set_terminator(struct.unpack('H', self.buf)[0])
        elif self.stage == 0:
            self.stage = 1
            self.set_terminator(2)
            peerlisten = struct.unpack('H', self.buf)[0]
            if not peerlisten:
                peerlisten = self.getpeername()[1] | 0x10000
            self.peer = core.addr_to_ident((self.getpeername()[0], peerlisten))
            self.hook = hooks.get(self.peer, None)
            try:
                old = nodes[self.peer]
                old.close_when_done()
                old.peer = None
            except KeyError:
                pass
            nodes[self.peer] = self
        self.buf = ''

    def handle_connect(self):
        pass

    def handle_close(self):
        nodes.pop(self.peer, None)
        self.close()

    def handle_expt(self):
        self.handle_close()

    def handle_error(self):
        self.handle_close()

        

class Listener(asyncore.dispatcher):
    def __init__(self, addr, backlog):
        asyncore.dispatcher.__init__(self)
        self.create_socket(socket.AF_INET, socket.SOCK_STREAM)
        self.bind(addr)
        self.listen(backlog)

    def handle_accept(self):
        (sock, remote) = self.accept()
        ch = Channel(sock)
        ch.init_accept()

    def writable(self):
        return False


def connect(ident):
    try:
        return nodes[ident]
    except KeyError:
        if ident.startswith('#'):
            return None
        ch = Channel()
        addr = core.ident_to_addr(ident)
        ch.init_connect(addr, port)
        return ch





class TimerItem:
    def __init__(self, timeout, callback):
        self.expire = time.time() + timeout
        self.callback = callback

    def __le__(self, other):
        return self.expire <= other.expire


def timer(timeout, hook, *args, **kwargs):
    heapq.heappush(timers, TimerItem(timeout, lambda: hook(*args, **kwargs)))


def init(addr, backlog):
    global port
    if addr:
        port = addr[1]
        Listener(addr, backlog)
    else:
        port = 0


def receive(ident, hook):
    if hook:
        hooks[ident] = hook
    else:
        hooks.pop(ident, None)
    try:
        nodes[ident].hook = hook
    except KeyError:
        pass


def send(msg, ident):
    ch = connect(ident)
    if ch:
        ch.push(struct.pack('H', len(msg)) + msg)


timers = []

def run():
    while True:
        asyncore.loop((timers[0].expire - time.time())
                      if timers else None)
        if not timers:
            break
        while timers and timers[0].expire < time.time():
            heapq.heappop(timers).callback()
